# -*- coding: utf-8 -*-

"""
Implementation of various merge strategies
"""

from __future__ import unicode_literals
from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

from typing import List, Dict, Any, AnyStr

from pyspark import RDD

from kgpipeline.util import import_direct_from_module

from kgpipeline.job import JobConfigError
from kgpipeline.lambdas import common_merge
from kgpipeline.lambdas import to_key_value, to_value_only, merge_override_with_delete_all
from kgpipeline.calc.util import count_by_key

from kgpipeline.util import IdUtil


def merge_fix_one(dict1, dict2):
    if dict1 is None:
        return dict2

    if dict2 is None:
        return dict1

    return merge_override_with_delete_all(dict1, dict2)


# a fix is a normal json stream text file
# it contains entities or links that need to be applied to current knowledge graph
# it has the following types:
# 1. only @id: it means entity/link with this id should be deleted
# 2. other: if @id exists in current kg, then do an override merge; else just add them
def merge_fix_rdd(rdd1, rdd2):
    rdd = to_key_value(rdd1)
    fix_rdd = to_key_value(rdd2)
    to_delete = fix_rdd.filter(lambda x: len(x[1]) == 1)
    rdd = rdd.subtractByKey(to_delete)
    return rdd.fullOuterJoin(fix_rdd.filter(lambda x: len(x[1]) > 1)).map(
        lambda x: merge_fix_one(*x[1]))


# noinspection PyUnusedLocal
def merge_fix(entity, link, entity_fix, link_fix,
              **kwargs):  # type: (RDD, RDD, RDD, RDD, Dict[AnyStr,Any]) -> List[RDD]
    if entity_fix and not entity_fix.isEmpty():
        entity = merge_fix_rdd(entity, entity_fix)

    if link_fix and not link_fix.isEmpty():
        link = merge_fix_rdd(link, link_fix)

    return [entity, link]


def substitute_id(row, key):
    if row[1][1] is not None:
        row[1][0][key] = row[1][1]
    return row[1][0]


def rewrite_link(link, id_map_table):
    return link.map(lambda x: (x["in"], x)).leftOuterJoin(id_map_table).map(lambda x: substitute_id(x, "in")).map(
        lambda x: (x["out"], x)
    ).leftOuterJoin(id_map_table).map(lambda x: substitute_id(x, "out"))


def has_all_keys(dict1, keys):
    for key in keys:
        if key not in dict1:
            return False
    return True


def get_id_by_keys(dict1, keys):
    return IdUtil.any2hash([dict1[k] for k in keys])


def merge_by_type_and_key_one_type(e1, e2, link2, type_name, key_names, merge_func):
    entity_of_type_1 = e1.filter(lambda x: has_all_keys(x[1], key_names) and type_name in x[1]["@type"]).map(
        lambda x: (get_id_by_keys(x[1], key_names), x[1]))
    entity_of_type_2 = e2.filter(lambda x: has_all_keys(x[1], key_names) and type_name in x[1]["@type"]).map(
        lambda x: (get_id_by_keys(x[1], key_names), x[1]))

    # checking validity of keys (non repeating)
    if count_by_key(entity_of_type_1).filter(lambda x: x[1] > 1).count() > 0:
        raise Exception("There are repeating IDs for selected keys: {0}!".format(" ".join(key_names)))

    if count_by_key(entity_of_type_2).filter(lambda x: x[1] > 1).count() > 0:
        raise Exception("There are repeating IDs for selected keys: {0}!".format(" ".join(key_names)))

    matched = entity_of_type_1.join(entity_of_type_2).map(lambda x: (x[1][0], x[1][1]))

    # @id1 => entity2
    entity_to_merge = matched.map(lambda x: (x[0]["@id"], x[1]))

    # @id2 => entity2
    entity_to_subtract = matched.map(lambda x: (x[1]["@id"], x[1]))

    # @id2 => @id1
    link_id_map = matched.map(lambda x: (x[1]["@id"], x[0]["@id"])).filter(lambda x: x[0] != x[1])

    e1 = e1.leftOuterJoin(entity_to_merge).mapValues(lambda x: merge_func(*x))
    e2 = e2.subtractByKey(entity_to_subtract)

    if not link_id_map.isEmpty():
        link2 = rewrite_link(link2, link_id_map)

    return e1, e2, link2


def merge_by_type_and_key(entity1, link1, entity2, link2, **kwargs):  # type: (RDD, RDD, RDD, RDD, Dict) -> List[RDD]
    """
    把外部entity和link融入到当前图谱，外部entity和link是和当前同一个schema
    使用类型和简单字段来确定是否是同一条实体, 若是那么进行融合并对链接外部链接进行修改，指向融合后的当前图谱的实体
    """
    verbose = "verbose" in kwargs

    if "merge_type_key_list" not in kwargs:
        raise JobConfigError("it is mandatory to set type_and_key to list of tuple: containing the type and key to use")

    merge_rest = kwargs.get("merge_rest", False)

    # [
    #     {"type": "MutualFund", "merge_keys": ["identifier"]}
    # ]
    merge_type_key = kwargs["merge_type_key_list"]

    reducer_name = kwargs.get("reducer")  # type: AnyStr

    if reducer_name:
        reducer = import_direct_from_module(reducer_name)
    else:
        reducer = common_merge

    e1 = entity1.map(lambda x: (x["@id"], x))
    e2 = entity2.map(lambda x: (x["@id"], x))

    if verbose:
        print("Initial: e1: {0}, e2: {1}, link2: {2}".format(e1.count(), e2.count(), link2.count()))

    for type_keys in merge_type_key:
        type_name = type_keys["type"]
        key_names = type_keys["merge_keys"]
        e1, e2, link2 = merge_by_type_and_key_one_type(e1, e2, link2, type_name, key_names, reducer)
        if verbose:
            print("Type: {0}, Key: {1}".format(type_name, key_names))
            print("\te1: {0}, e2: {1}, link2: {2}".format(e1.count(), e2.count(), link2.count()))

    if merge_rest:
        return [to_value_only(e1.union(e2)), link1.union(link2)]
    else:
        return [to_value_only(e1), link1]


# noinspection PyUnusedLocal
def simple_merge(entity1, link1, entity2, link2, **kwargs):  # type: (RDD, RDD, RDD, RDD, Dict[AnyStr,Any]) -> List[RDD]
    """
    简单融合：根据@id来融合两个图谱，这个和默认的production job的reduce步骤实现了一样的逻辑
    """
    reducer_name = kwargs.get("reducer")  # type: AnyStr

    if reducer_name:
        reducer = import_direct_from_module(reducer_name)
    else:
        reducer = common_merge

    if entity1 is None or entity1.isEmpty():
        e = entity2
    elif entity2 is None or entity2.isEmpty():
        e = entity1
    else:
        e1 = to_key_value(entity1)
        e2 = to_key_value(entity2)
        e = to_value_only(e1.union(e2).reduceByKey(reducer))

    if link1 is None or link1.isEmpty():
        r = link2
    elif entity2 is None or link2.isEmpty():
        r = link1
    else:
        l1 = to_key_value(link1)
        l2 = to_key_value(link2)
        r = to_value_only(l1.union(l2).reduceByKey(reducer))

    return [e, r]
